#include "stdafx.h"
#include <string>
#include <vector>
#define COMPILING_DYNLINK
#include "dynlink.h"

#include <tchar.h>
#include <strsafe.h>

static std::wstring a2w(const std::string &s_in)
{
	std::wstring s_out;
	int count = MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, s_in.c_str(), -1, NULL, 0);
	if(count==0)
		return s_out;
	 
	void* pBuffer = (void*) new wchar_t[count];
	int result = MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, s_in.c_str(), -1, (LPWSTR)pBuffer, count);
	if (result > 0)
		s_out = (wchar_t *) pBuffer;

	delete [] pBuffer;
	return s_out;
}    

static std::string w2a(const std::wstring &s_in)
{ 
	std::string s_out;
	int count = WideCharToMultiByte(CP_ACP, 0, s_in.c_str(), -1, NULL, 0, NULL, NULL);
	if(count==0)
		return s_out;
	 
	void* pBuffer = (void*) new char[count];
	int result = WideCharToMultiByte(CP_ACP, 0, s_in.c_str(), -1, (LPSTR)pBuffer, count, NULL, NULL);
	if (result > 0)
		s_out = (char *) pBuffer;

	delete [] pBuffer;
	return s_out;
}

void CDynlinkBase::get_error(DWORD dwError, std::string & s_err) 
{ 
	// Retrieve the system error message for the last-error code
	LPVOID lpMsgBuf;
	LPVOID lpDisplayBuf;
	
	FormatMessageA(
		FORMAT_MESSAGE_ALLOCATE_BUFFER | 
		FORMAT_MESSAGE_FROM_SYSTEM |
		FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL,
		dwError,
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPSTR) &lpMsgBuf,
		0, NULL );

	// Display the error message and exit the process
	size_t len = lstrlenA((LPCSTR)lpMsgBuf);
	((LPSTR)lpMsgBuf)[len-1] = 0;
	lpDisplayBuf = (LPVOID)LocalAlloc(LMEM_ZEROINIT, (len +  40) * sizeof(CHAR)); 
	StringCchPrintfA((LPSTR)lpDisplayBuf, LocalSize(lpDisplayBuf) / sizeof(CHAR), "%d: %s", dwError, lpMsgBuf); 
	s_err = (LPCSTR) lpDisplayBuf;
	LocalFree(lpMsgBuf);
	LocalFree(lpDisplayBuf);
}

void CDynlinkBase::get_error(DWORD dwError, std::wstring & s_err) 
{ 
	// Retrieve the system error message for the last-error code
	LPVOID lpMsgBuf;
	LPVOID lpDisplayBuf;
	
	FormatMessageW(
		FORMAT_MESSAGE_ALLOCATE_BUFFER | 
		FORMAT_MESSAGE_FROM_SYSTEM |
		FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL,
		dwError,
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPWSTR) &lpMsgBuf,
		0, NULL );

	// Display the error message and exit the process
	size_t len = lstrlenW((LPCWSTR)lpMsgBuf);
	((LPWSTR)lpMsgBuf)[len-1] = 0;
	lpDisplayBuf = (LPVOID)LocalAlloc(LMEM_ZEROINIT, (len +  40) * sizeof(WCHAR)); 
	StringCchPrintfW((LPWSTR)lpDisplayBuf, LocalSize(lpDisplayBuf) / sizeof(WCHAR), L"%d: %s", dwError, lpMsgBuf); 
	s_err = (LPCWSTR) lpDisplayBuf;
	LocalFree(lpMsgBuf);
	LocalFree(lpDisplayBuf);
}

DWORD CDynlinkBase::get_error(std::wstring &s_err) const
{
	DWORD dwErr = get_error();
	get_error(dwErr, s_err);
	return dwErr;
}

DWORD CDynlinkBase::get_error(std::string &s_err) const
{
	DWORD dwErr = get_error();
	get_error(dwErr, s_err);
	return dwErr;
}

CDynlinkLib::CDynlinkLib(LPCWSTR sName, CDynlinkLib::mode d_mode)
{
	// Will crash in case of a null pointer, this is by design at the moment
	// as we insist on a fall-back library name.
	// This name with a ".dll" added to it is the default name under
	// which the library will be loaded.
	init(sName, d_mode);
}

CDynlinkLib::CDynlinkLib(LPCSTR sName, CDynlinkLib::mode d_mode)
{
	// Will crash in case there is a NULL pointer
	init(a2w(sName), d_mode);
}

void CDynlinkLib::init(const std::wstring& sName, mode d_mode)
{
	m_d_throw = 0;
	m_b_auto_load = true;
	m_mode = d_mode;
	m_h_dll = NULL;
	m_d_loaded = DYNLINK_ERR_NOTLOADED;
	m_s_name_bak = sName;
	m_s_name = sName;
}

int CDynlinkLib::load(LPCSTR s_name, bool b_retry)
{
	// Need to pass on a NULL in case we have no name
	return load(s_name ? a2w(s_name).c_str() : 0, b_retry);
}

DWORD CDynlinkSymbol::get_error() const
{
	if (is_loaded())
	{
		// m_dwError == ERROR_SUCCESS should always be true here
		return ERROR_SUCCESS;
	}
	else
	{
		// Some call to GetProcAddress() has obviously occurred
		if (m_dwError != ERROR_SUCCESS)
			return m_dwError;
		else
			return m_p_lib ->get_error();
	}
}

DWORD CDynlinkLib::get_error() const
{
	return m_dwError;
}


void * CDynlinkSymbol::check()
{
	// Auto load if this symbol has not yet been unsuccessfully loaded
	if ((! m_p_ptr) && (m_dwError == ERROR_SUCCESS) && (m_p_lib ->m_b_auto_load)) 
		m_p_lib ->load();

	switch(m_p_lib ->m_d_throw)
	{
	// Throw if we are not loaded or if we are a stub
	case 2:
		if (!is_loaded())
			throw this;
	// Throw if we are not loaded 
	case 1:
		if (!is_loaded_or_stub())
			throw this;
	// Don't do anything
	default:
		;
	}
	// Return the pointer of the stub, i.e. this behaves as if check()
	// was not called
	return m_p_ptr;
}


int CDynlinkLib::load_stubs()
{
	free();

	switch(m_mode)
	{
	default:
	case ex:
	case imp:
		return (m_d_loaded = DYNLINK_ERR_INVALID_METHOD); 
	case exstub:
	{
		bool b_all = true;
		m_h_dll = (HMODULE)-1;

		std::vector<CDynlinkSymbol *>::iterator it = m_vecSymbols.begin();
		for(; it != m_vecSymbols.end(); it++)
			b_all &= (*it) ->load(m_h_dll);			
			
		// b_all should always be true as if one function has a stub, so 
		// should the others. But if someone is not using the macros ...
		return (m_d_loaded = b_all ? DYNLINK_SUCCESS : DYNLINK_ERR_NULLPTRS);
	}
	};
}
		

int CDynlinkLib::load(LPCWSTR s_name_in, bool b_retry)
{
	bool b_stubs = false;
	// If no name is given we always use the name that was set during
	// creation, i.e. the name that was defined in the header file if the
	// library is used correctly.
	std::wstring s_name = s_name_in ? s_name_in : (m_s_name_bak + L".dll");

	// Force a retry if the name has changed since the last try.
	if (s_name != m_s_name)
	{
		m_s_name = s_name;
		b_retry = true;
	}
		
	// If we already tried either retry or clear the old state and retry from
	// the start
	if (m_d_loaded != DYNLINK_ERR_NOTLOADED)
	{
		if (! b_retry)
			return m_d_loaded;
		else
			free();
	}
		
	m_dwError = ERROR_SUCCESS;

	switch(m_mode)
	{
	case imp:
		m_s_name = m_s_name_bak;
		m_h_dll = GetModuleHandleW(m_s_name_bak.c_str());
		m_d_loaded = DYNLINK_SUCCESS;
		if (m_s_name_bak != s_name)
			return DYNLINK_ERR_INVALID_METHOD;
		else
			return DYNLINK_SUCCESS;
	case exstub:
		// Load explicitely with stubs
		b_stubs = true;
	case ex:
		// without stubs
		m_h_dll = LoadLibraryW(m_s_name.c_str());
		if (m_h_dll)
		{
			bool b_all = true;
			std::vector<CDynlinkSymbol *>::iterator it = m_vecSymbols.begin();
			for(; it != m_vecSymbols.end(); it++)
				b_all &= (*it) ->load(m_h_dll);			
				
			if (! b_all)
				m_d_loaded = b_stubs ? DYNLINK_ERR_STUBS : DYNLINK_ERR_NULLPTRS;
			else
				m_d_loaded = DYNLINK_SUCCESS;
		}
		else
		{
			m_d_loaded = DYNLINK_ERR_LOADLIBRARY;
			m_dwError = GetLastError();
		}
		return m_d_loaded;

	default:
		return m_d_loaded = DYNLINK_ERR_INVALID_METHOD;
	}
}

void CDynlinkLib::add(CDynlinkSymbol * p_symbol)
{
	m_vecSymbols.push_back(p_symbol);
}

CDynlinkSymbol::CDynlinkSymbol(CDynlinkLib &lib, LPCSTR s_name)
{
	m_s_name = s_name;
	m_p_stub = NULL;
	m_p_ptr  = NULL;

	m_dwError = ERROR_SUCCESS;
	m_p_lib = &lib;
	lib.add(this);
}

bool CDynlinkSymbol::load(HMODULE hDll)
{
	// m_p_ptr should always be NULL here
	if (hDll == ((HMODULE)-1))
	{
		// Load the stub and fake that it is not the stub. 
		if (m_p_stub)
		{
			m_p_ptr = m_p_stub;
			m_p_stub = ((size_t *) m_p_stub - 1);
			m_dwError = ERROR_SUCCESS;
			return true;
		}
		else
		{
			m_dwError = ERROR_PROC_NOT_FOUND;
			return false;
		}
	}

	if (! (m_p_ptr = GetProcAddress(hDll, m_s_name.c_str())))
	{
		// Set the error, place the stub and return failure
		m_dwError = GetLastError();
		m_p_ptr = m_p_stub;
		return false;
	}
	else
	{
		m_dwError = ERROR_SUCCESS;
		return true;
	}

}	

void CDynlinkSymbol::free()
{
	m_dwError = ERROR_SUCCESS;
	// Restore the stub if it was faking a real load
	if ((m_p_ptr == (size_t *) m_p_stub + 1))
		m_p_stub = (size_t *) m_p_stub + 1;
	m_p_ptr = NULL;
}

bool CDynlinkLib::free()
{
	if (! m_h_dll)
		return TRUE;
	
	switch(m_mode)
	{
	case imp:
		// Cannot unload
		return false;
	case ex:
	case exstub:
		if ((m_h_dll == (HMODULE)-1) || FreeLibrary(m_h_dll))
		{
			std::vector<CDynlinkSymbol *>::iterator it = m_vecSymbols.begin();
			for(; it != m_vecSymbols.end(); it++)
				(*it) ->free();
			break;
		}
		else
		{
			return FALSE;
		}
	default:
		return false;
	}
	m_h_dll = NULL;
	m_d_loaded = DYNLINK_ERR_NOTLOADED;
	
	return true;
}



bool CDynlinkLib::get_filename(std::string &s_fn) const
{
	std::wstring s_fnw;
	bool bRet = get_filename(s_fnw);

	s_fn = w2a(s_fnw);
	return bRet;
}

bool CDynlinkLib::get_filename(std::wstring &s_fn) const
{
	WCHAR fn[1024];
	if (GetModuleFileNameW(get_module(), fn, 1024))
	{
		s_fn = fn;
		return true;
	}
	else
	{
		return false;
	}
}
